#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time       : 2025/3/3 下午4:59
@Auther     : 635304837@qq.com
@File       : chatbackup.py
"""
import json
import os
import time
import logging
from ctypes import CDLL, c_char_p, c_int, c_longlong, c_uint, c_void_p, Structure, byref

# 日志配置
logging.basicConfig(
    filename="sdktools.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)


# 定义 C 结构体
class Slice_t(Structure):
    _fields_ = [
        ("buf", c_char_p),
        ("len", c_int)
    ]

class MediaData_t(Structure):
    _fields_ = [
        ("outindexbuf", c_char_p),
        ("out_len", c_int),
        ("data", c_char_p),
        ("data_len", c_int),
        ("is_finish", c_int)
    ]

# 加载共享库
lib = CDLL("./libWeWorkFinanceSdk_C.so")

# 定义函数签名
lib.NewSdk.restype = c_void_p
lib.Init.argtypes = [c_void_p, c_char_p, c_char_p]
lib.Init.restype = c_int
lib.GetChatData.argtypes = [c_void_p, c_longlong, c_uint, c_char_p, c_char_p, c_int, c_void_p]
lib.GetChatData.restype = c_int
lib.GetMediaData.argtypes = [c_void_p, c_char_p, c_char_p, c_char_p, c_char_p, c_int, c_void_p]
lib.GetMediaData.restype = c_int
lib.DecryptData.argtypes = [c_char_p, c_char_p, c_void_p]
lib.DecryptData.restype = c_int
lib.DestroySdk.argtypes = [c_void_p]
lib.NewSlice.restype = c_void_p
lib.FreeSlice.argtypes = [c_void_p]
lib.NewMediaData.restype = c_void_p
lib.FreeMediaData.argtypes = [c_void_p]

# 配置类
class Config:
    def __init__(self):
        self.corpid = "ww*************"
        self.secret = "y**********************4"
        self.proxy = ""
        self.passwd = ""
        self.timeout = 5
        self.limit = 100
        self.save_dir = "saved_data"
        self.current_seq = 0

# SDK 类
class WeWorkFinanceSdk:
    def __init__(self, config):
        self.config = config
        self.sdk = lib.NewSdk()
        ret = lib.Init(self.sdk, config.corpid.encode(), config.secret.encode())
        if ret != 0:
            logging.error(f"Init SDK failed, error code: {ret}")
            raise Exception(f"Init SDK failed, error code: {ret}")

    def __del__(self):
        lib.DestroySdk(self.sdk)

    def get_chat_data(self):
        slice_ptr = lib.NewSlice()
        ret = lib.GetChatData(
            self.sdk,
            self.config.current_seq,
            self.config.limit,
            self.config.proxy.encode(),
            self.config.passwd.encode(),
            self.config.timeout,
            slice_ptr
        )
        if ret != 0:
            lib.FreeSlice(slice_ptr)
            logging.error(f"GetChatData failed, error code: {ret}")
            raise Exception(f"GetChatData failed, error code: {ret}")

        slice = Slice_t.from_address(slice_ptr)
        data = slice.buf[:slice.len].decode('utf-8')
        lib.FreeSlice(slice_ptr)
        return json.loads(data)

    def get_media_data(self, sdk_file_id, save_file):
        index = ""
        isfinish = 0
        media_data_list = []
        while isfinish == 0:
            media_data_ptr = lib.NewMediaData()
            ret = lib.GetMediaData(
                self.sdk,
                index.encode(),
                sdk_file_id.encode(),
                self.config.proxy.encode(),
                self.config.passwd.encode(),
                self.config.timeout,
                media_data_ptr
            )
            if ret != 0:
                lib.FreeMediaData(media_data_ptr)
                logging.error(f"GetMediaData failed, error code: {ret}")
                raise Exception(f"GetMediaData failed, error code: {ret}")

            media_data = MediaData_t.from_address(media_data_ptr)
            media_data_dict = {
                "outindexbuf": media_data.outindexbuf.decode(),
                "data": media_data.data[:media_data.data_len].tobytes(),
                "is_finish": media_data.is_finish
            }
            media_data_list.append(media_data_dict)

            index = media_data.outindexbuf.decode()
            isfinish = media_data.is_finish
            lib.FreeMediaData(media_data_ptr)

        # 保存媒体数据
        with open(save_file, "wb") as fp:
            for data in media_data_list:
                fp.write(data["data"])

        logging.info(f"MediaData saved to {save_file}")
        return save_file

    def decrypt_data(self, encrypt_key, encrypt_msg):
        slice_ptr = lib.NewSlice()
        ret = lib.DecryptData(
            encrypt_key.encode(),
            encrypt_msg.encode(),
            slice_ptr
        )
        if ret != 0:
            lib.FreeSlice(slice_ptr)
            logging.error(f"DecryptData failed, error code: {ret}")
            raise Exception(f"DecryptData failed, error code: {ret}")

        slice = Slice_t.from_address(slice_ptr)
        data = slice.buf[:slice.len].decode('utf-8')
        lib.FreeSlice(slice_ptr)
        return json.loads(data)

    def parse_and_save_message(self, message):
        try:
            msg_type = message.get("type", "text")
            content = message.get("content", "")
            file_id = message.get("file_id", None)

            save_dir = os.path.join(self.config.save_dir, msg_type)
            os.makedirs(save_dir, exist_ok=True)

            if msg_type == "text":
                save_file = os.path.join(save_dir, f"{int(time.time())}.txt")
                with open(save_file, "w") as fp:
                    fp.write(content)
            elif msg_type in ["image", "video", "audio", "file"]:
                if file_id:
                    save_file = os.path.join(save_dir, file_id.split('/')[-1])
                    self.get_media_data(file_id, save_file)
                else:
                    logging.warning(f"Media type {msg_type} without file_id")
                    save_file = None
            else:
                logging.warning(f"Unknown message type: {msg_type}")
                save_file = None

            return save_file
        except Exception as e:
            logging.error(f"Failed to parse message: {str(e)}")
            return None

    def save_chat_records_continuously(self):
        exit_loop = False  # 标记是否退出循环
        while not exit_loop:
            try:
                response = self.get_chat_data()
                if response.get("errcode") != 0:
                    logging.error(f"GetChatData error: {response.get('errmsg')}")
                    exit_loop = True
                    continue

                chat_records = response.get("chatdata", [])
                if not chat_records:
                    logging.info("No more chat records to fetch")
                    exit_loop = True
                    continue

                # 保存最大 seq
                max_seq = max(record.get("seq", 0) for record in chat_records)
                self.config.current_seq = max_seq + 1  # 更新下次拉取的 seq

                # 处理聊天记录
                for record in chat_records:
                    msgid = record.get("msgid")
                    encrypt_key = record.get("encrypt_random_key")
                    encrypt_msg = record.get("encrypt_chat_msg")

                    try:
                        decrypted_msg = self.decrypt_data(encrypt_key, encrypt_msg)
                        logging.info(f"Decrypted message {msgid}: {decrypted_msg}")
                    except Exception as e:
                        logging.error(f"Failed to decrypt message {msgid}: {str(e)}")
                        continue

                    try:
                        save_path = self.parse_and_save_message(decrypted_msg)
                        if save_path:
                            logging.info(f"Message {msgid} saved to {save_path}")
                    except Exception as e:
                        logging.error(f"Failed to save message {msgid}: {str(e)}")

                # 延迟一段时间再拉取下一批数据
                time.sleep(60)

            except Exception as e:
                logging.error(f"SDK Error: {str(e)}")
                exit_loop = True

# 配置
config = Config()
config.save_dir = "saved_data"

# 运行
if __name__ == "__main__":
    try:
        sdk = WeWorkFinanceSdk(config)
        sdk.save_chat_records_continuously()
    except Exception as e:
        logging.error(f"SDK Error: {str(e)}")